#!/usr/bin/python
# -*- coding: utf8 -*-
#-------------------------------------------------------------------------------
# Name:        模块1
# Purpose:
#
# Author:      admin
#
# Created:     13/05/2018
# Copyright:   (c) admin 2018
# Licence:     <your licence>
#-------------------------------------------------------------------------------
from __future__ import print_function
import sqlite3,os,sys
from multiprocessing import Process
import threading
import socket
import subprocess
import traceback,time
encode='gb2312'
def read_stdout(local_conn, remote_ip,cmd, pread = None):
    if not pread:
        pread = cmd.stdout
    s = ''
    while 1:
        try:
##            if not cmd.stdout.readable():
##                time.sleep(1)
##                continue
            if cmd.poll() is not None :
##                print('read_stdout exit')
                break
            if hasattr(pread, 'read'):
                s = pread.read(1)
            else:
                s = os.read(pread, 1)
##            print('read_stdout ' + str(cmd.returncode ))
            print(s,end='')
            local_conn.send(s)
            #send(local_conn,s)
        except KeyboardInterrupt as e:
            break
        except:
            print(s)
            traceback.print_exc()
            break

def bytes2str(s,encode):
    if sys.version_info[0] == 2:
        return s
    return str(s,encode)

def str2bytes(b,encode):
    if sys.version_info[0] == 2:
        return b
    return bytes(b, encode)

def tcp_mapping_request(local_conn, remote_ip):
    enable_cmd = 1
    if enable_cmd:
        pread,pwrite = os.pipe()
        cmd = subprocess.Popen(r'cmd',shell=True,stdout=pwrite,stderr=pwrite, stdin = subprocess.PIPE)
        t = threading.Thread(target=read_stdout, args=(local_conn, remote_ip, cmd, pread), )
        t.daemon = True
        t.start()
        os.close(pwrite)
    while True:
        try:
            if cmd.poll() is not None:
                print('end %s %d'%(remote_ip[0], cmd.returncode))
                break
            try:
                from_recv = local_conn.recv(8096)
            except:
                continue
            #from_recv = recv(local_conn, 8096)
            if len(from_recv) == 0:continue
            msg =bytes2str(from_recv,encode)
            print ('%s %s' % (remote_ip[0],msg))
            if enable_cmd:
                cmd.stdin.write(from_recv)
                cmd.stdin.write(str2bytes('\r\n','utf8'))
                cmd.stdin.flush()
            else:
                cmd = subprocess.Popen(msg,shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE)#, stdin = subprocess.PIPE)
                s = cmd.stdout.read()
                if not s:
                    s = cmd.stderr.read()
                local_conn.send(s)
        except KeyboardInterrupt as e:
            break
##        except ConnectionResetError as e:
##            break
        except Exception:
          traceback.print_exc()
    cmd.kill()
    send(local_conn, 'exit')
    local_conn.close()
    if enable_cmd:
        os.close(pread)

def server():
    sk_obj = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    sk_obj.settimeout(3)
    sk_obj.bind(('0.0.0.0',8000))
    sk_obj.listen(5)
    while True:
        try:
            conn,ipaddr = sk_obj.accept()
        except KeyboardInterrupt as e:
            break
        except socket.timeout as e:
            continue
        print ('connection from ip: %s' % ipaddr[0])
        t = threading.Thread(target=tcp_mapping_request, args=(conn, ipaddr), )
        t.daemon = True
        t.start()
    sk_obj.shutdown(socket.SHUT_RDWR)
    sk_obj.close()

def read_socket(sk_obj):
    while 1:
        data = recv(sk_obj, 1)
        #print ('%s' % data.decode('utf-8'))
        print(data,end='')

def client():
    sk_obj=socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    sk_obj.settimeout(3)
    sk_obj.connect(('10.185.215.167',8000))
##    threading.Thread(target=read_socket, args=(sk_obj, ), daemon=True).start()
    def recv_data(sk_obj):
        s = ''
        while 1:
            data = recv(sk_obj, 1)
            s+=data
            #print ('%s' % data.decode('utf-8'))
            print(data,end='')
            data = s.splitlines()[-1]
            if data.endswith('>') and ':\\' in data:
                break
    recv_data(sk_obj)
    while True:
        msg = raw_input()
        if not msg:continue
        send(sk_obj, msg)
##        if msg == 'exit':
##            break
        recv_data(sk_obj)
    sk_obj.close()

def recv(skt,cnt):
    if sys.version_info[0] == 2:
        return skt.recv(cnt)

    s = bytes()
    while 1:
        s += skt.recv(cnt)
        try:
            return str(s, encode)
        except UnicodeDecodeError as e:
            continue

def send(skt,s):
    if sys.version_info[0] == 2:
        return skt.send(s)
    return skt.send(bytes(s,encode))

def test_stdin():
    of = open('d:\\a.log','w')
    while 1:
        s = input('cmd:')
        of.write(s)
        of.write('\n')
        if s == 'exit':
            break
    of.close()

def main():
    if '-s' in sys.argv:
        return server()
    elif '-i' in sys.argv:
        test_stdin()
    else:
        return client()


if __name__ == '__main__':
    main()